package org.gensokyo.crawler.util

import java.io.IOException
import java.nio.charset.Charset

import com.google.common.collect.Sets
import org.apache.commons.io.IOUtils
import org.apache.commons.lang3.StringUtils
import org.apache.http.client.config.{CookieSpecs, RequestConfig}
import org.apache.http.client.methods.{CloseableHttpResponse, HttpUriRequest, RequestBuilder}
import org.apache.http.impl.client.CloseableHttpClient
import org.apache.http.util.EntityUtils
import org.apache.http.{HttpHost, HttpResponse, NameValuePair}
import org.jsoup.Jsoup
import org.slf4j.LoggerFactory
import us.codecraft.webmagic.downloader.{AbstractDownloader, HttpClientGenerator}
import us.codecraft.webmagic.selector.PlainText
import us.codecraft.webmagic.utils.{HttpConstant, UrlUtils}
import us.codecraft.webmagic.{Page, Request, Site, Task}

/**
  * Created by liusen on 16-12-1.
  */
class ScalaHttpClientDownloader extends AbstractDownloader {

  private val logger = LoggerFactory.getLogger(getClass)
  private val httpClients: java.util.Map[String, CloseableHttpClient] = new java.util.HashMap[String, CloseableHttpClient]

  private val httpClientGenerator: HttpClientGenerator = new HttpClientGenerator

  private def getHttpClient(site: Site): CloseableHttpClient = {
    if (site == null) return httpClientGenerator.getClient(null)
    val domain: String = site.getDomain
    var httpClient: CloseableHttpClient = httpClients.get(domain)
    if (httpClient == null) this synchronized {
      httpClient = httpClients.get(domain)
      if (httpClient == null) {
        httpClient = httpClientGenerator.getClient(site)
        httpClients.put(domain, httpClient)
      }
    }
    httpClient
  }

  def download(request: Request, task: Task): Page = {
    var site: Site = null
    if (task != null) site = task.getSite
    var acceptStatCode: java.util.Set[Integer] = null
    var charset: String = null
    var headers: java.util.Map[String, String] = null
    if (site != null) {
      acceptStatCode = site.getAcceptStatCode
      charset = site.getCharset
      headers = site.getHeaders
    }
    else acceptStatCode = Sets.newHashSet(200)
    logger.info("downloading page {}", request.getUrl)
    var httpResponse: CloseableHttpResponse = null
    var statusCode: Int = 0
    try {
      val httpUriRequest: HttpUriRequest = getHttpUriRequest(request, site, headers)
      httpResponse = getHttpClient(site).execute(httpUriRequest)
      statusCode = httpResponse.getStatusLine.getStatusCode
      request.putExtra(Request.STATUS_CODE, statusCode)
      if (statusAccept(acceptStatCode, statusCode)) {
        val page: Page = handleResponse(request, charset, httpResponse, task)
        onSuccess(request)
        page
      }
      else {
        logger.warn("code error " + statusCode + "\t" + request.getUrl)
        null
      }
    } catch {
      case e: IOException => {
        logger.warn("download page " + request.getUrl + " error", e)
        if (site.getCycleRetryTimes > 0) return addToCycleRetry(request, site)
        onError(request)
        null
      }
    } finally {
      request.putExtra(Request.STATUS_CODE, statusCode)
      try
          if (httpResponse != null) {
            //ensure the connection is released back to pool
            EntityUtils.consume(httpResponse.getEntity)
          }

      catch {
        case e: IOException => {
          logger.warn("close response fail", e)
        }
      }
    }
  }

  def setThread(thread: Int) {
    httpClientGenerator.setPoolSize(thread)
  }

  protected def statusAccept(acceptStatCode: java.util.Set[Integer], statusCode: Int): Boolean = acceptStatCode.contains(statusCode)

  protected def getHttpUriRequest(request: Request, site: Site, headers: java.util.Map[String, String]): HttpUriRequest = {
    val requestBuilder: RequestBuilder = selectRequestMethod(request).setUri(request.getUrl)
    if (headers != null) {
      import scala.collection.JavaConversions._
      for (headerEntry <- headers.entrySet) {
        requestBuilder.addHeader(headerEntry.getKey, headerEntry.getValue)
      }
    }
    val requestConfigBuilder: RequestConfig.Builder = RequestConfig.custom.setConnectionRequestTimeout(site.getTimeOut).setSocketTimeout(site.getTimeOut).setConnectTimeout(site.getTimeOut).setCookieSpec(CookieSpecs.BEST_MATCH)
    if (site.getHttpProxyPool != null && site.getHttpProxyPool.isEnable) {
      val host: HttpHost = site.getHttpProxyFromPool
      requestConfigBuilder.setProxy(host)
      request.putExtra(Request.PROXY, host)
    }
    else if (site.getHttpProxy != null) {
      val host: HttpHost = site.getHttpProxy
      requestConfigBuilder.setProxy(host)
      request.putExtra(Request.PROXY, host)
    }
    requestBuilder.setConfig(requestConfigBuilder.build)
    requestBuilder.build
  }

  protected def selectRequestMethod(request: Request): RequestBuilder = {
    val method: String = request.getMethod
    if (method == null || method.equalsIgnoreCase(HttpConstant.Method.GET)) {
      //default get
      return RequestBuilder.get
    }
    else if (method.equalsIgnoreCase(HttpConstant.Method.POST)) {
      val requestBuilder: RequestBuilder = RequestBuilder.post
      val nameValuePair: Array[NameValuePair] = request.getExtra("nameValuePair").asInstanceOf[Array[NameValuePair]]
      if (nameValuePair != null && nameValuePair.length > 0) {
        //        requestBuilder.setEntity(new UrlEncodedFormEntity(util.Arrays.asList(nameValuePair), "UTF-8"))
        //nameValuePair.foreach(requestBuilder.addParameter)
      }
      return requestBuilder
    }
    else if (method.equalsIgnoreCase(HttpConstant.Method.HEAD)) return RequestBuilder.head
    else if (method.equalsIgnoreCase(HttpConstant.Method.PUT)) return RequestBuilder.put
    else if (method.equalsIgnoreCase(HttpConstant.Method.DELETE)) return RequestBuilder.delete
    else if (method.equalsIgnoreCase(HttpConstant.Method.TRACE)) return RequestBuilder.trace
    throw new IllegalArgumentException("Illegal HTTP Method " + method)
  }

  @throws[IOException]
  protected def handleResponse(request: Request, charset: String, httpResponse: HttpResponse, task: Task): Page = {
    val content: String = getContent(charset, httpResponse)
    val page: Page = new Page
    page.setRawText(content)
    page.setUrl(new PlainText(request.getUrl))
    page.setRequest(request)
    page.setStatusCode(httpResponse.getStatusLine.getStatusCode)
    page
  }

  @throws[IOException]
  protected def getContent(charset: String, httpResponse: HttpResponse): String = if (charset == null) {
    val contentBytes: Array[Byte] = IOUtils.toByteArray(httpResponse.getEntity.getContent)
    val htmlCharset: String = getHtmlCharset(httpResponse, contentBytes)
    if (htmlCharset != null) new String(contentBytes, htmlCharset)
    else {
      logger.warn("Charset autodetect failed, use {} as charset. Please specify charset in Site.setCharset()", Charset.defaultCharset)
      new String(contentBytes)
    }
  }
  else IOUtils.toString(httpResponse.getEntity.getContent, charset)

  @throws[IOException]
  protected def getHtmlCharset(httpResponse: HttpResponse, contentBytes: Array[Byte]): String = {
    var charset: String = null
    // charset
    // 1、encoding in http header Content-Type
    val value: String = httpResponse.getEntity.getContentType.getValue
    charset = UrlUtils.getCharset(value)
    if (StringUtils.isNotBlank(charset)) {
      logger.debug("Auto get charset: {}", charset)
      return charset
    }
    // use default charset to decode first time
    val defaultCharset: Charset = Charset.defaultCharset
    val content: String = new String(contentBytes, defaultCharset.name)
    // 2、charset in meta
    if (StringUtils.isNotEmpty(content)) {
      val document = Jsoup.parse(content)
      val links = document.select("meta")
      import scala.collection.JavaConversions._
      for (link <- links) {
        // 2.1、html4.01 <meta http-equiv="Content-Type" content="text/html; charset=UTF-8" />
        var metaContent: String = link.attr("content")
        val metaCharset: String = link.attr("charset")
        if (metaContent.indexOf("charset") != -1) {
          metaContent = metaContent.substring(metaContent.indexOf("charset"), metaContent.length)
          charset = metaContent.split("=")(1)
        } else if (StringUtils.isNotEmpty(metaCharset)) {
          charset = metaCharset
        }
      }
    }

    logger.debug("Auto get charset: {}", charset)
    // 3、todo use tools as cpdetector for content decode
    charset
  }

}
